# 导入库
import random
import numpy as np
import csv
import time


class GA:
    def __init__(self, fitness_func, n, population_size, generations, casei):
        # 外部传参
        self.fitness_func = fitness_func  # 适应度评估函数
        self.n = n  # 编码长度
        self.population_size = population_size  # 群体中个体数量
        self.generations = generations  # 迭代次数
        # 内部传参
        self.population = np.zeros((self.population_size, self.n))  # 存储当前种群的每个个体基因
        self.y_population = np.zeros(self.population_size)  # 存储当前种群每个个体适应度
        self.initialize_population()  # 初始化种群函数
        self.best_fitness = -np.inf  # 初始化最佳适应度为无穷小量
        self.number = 0  # 初始化迭代次数
        self.casei = casei

    # 初始化种群
    def initialize_population(self):
        # 使用numpy直接生成随机二进制数组
        self.population = np.random.randint(0, 2, (self.population_size, self.n))

    # 计算适应度并保存
    def cal_fitness(self):
        # 使用向量化操作计算适应度
        self.y_population = np.array([self.fitness_func(self.casei, ind) for ind in self.population])

        # 更新最佳适应度
        current_best = np.max(self.y_population)
        if current_best > self.best_fitness:
            self.best_fitness = current_best
            self.best_code = self.population[np.argmax(self.y_population)]

            # BestCode_result_file = "GAnew_code_" + str(self.casei) + ".csv"
            # with open(BestCode_result_file, 'w', encoding='utf-8', newline='') as f:
            #     write = csv.writer(f)
            #     write.writerow(list(self.best_code))

    # 选择操作，这里使用锦标赛选择
    def selection(self):
        # 使用numpy生成锦标赛选择的索引
        tournament_size = 3
        competitors = np.random.randint(0, self.population_size,
                                        (self.population_size, tournament_size))
        # 获取每组竞争者中的优胜者索引
        winners = np.argmax(self.y_population[competitors], axis=1)
        # 更新种群
        self.population = self.population[competitors[np.arange(len(winners)), winners]]

    # 交叉操作，这里使用单点交叉
    def crossover(self):
        # 生成所有交叉点
        crossover_points = np.random.randint(1, self.n, size=self.population_size // 2)
        # 获取配对的父代
        parents = self.population.reshape(-1, 2, self.n)

        # 创建子代数组
        children = np.zeros_like(parents)
        # 使用高级索引进行交叉操作
        for i, point in enumerate(crossover_points):
            children[i, 0] = np.concatenate([parents[i, 0, :point], parents[i, 1, point:]])
            children[i, 1] = np.concatenate([parents[i, 1, :point], parents[i, 0, point:]])

        # 更新种群
        self.population = children.reshape(self.population_size, self.n)

    # 变异操作，这里使用位翻转变异
    def mutation(self):
        # 生成变异掩码
        mutation_mask = np.random.random((self.population_size, self.n)) < 0.001
        # 应用变异
        self.population = np.where(mutation_mask, 1 - self.population, self.population)

    # 运行
    def run(self):
        for _ in range(self.generations):
            sti = time.time()
            self.number += 1  # 迭代次数加一
            self.cal_fitness()  # 评估适应度
            self.selection()  # 选择个体
            self.crossover()  # 交叉
            self.mutation()  # 变异

            Curve_result_file = "GAnew_result_" + str(self.casei) + ".csv"
            with open(Curve_result_file, 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                write.writerow([self.number * self.population_size, self.best_fitness.round(5)])
            print("Iter ", _, ": ", self.best_fitness, "Running time: ", time.time() - sti, 's')

        best_code = self.population[np.argmax(self.y_population)]  # 输出适应度最大对应的最优编码
        np.set_printoptions(threshold=np.inf)  # 显示数组的全部元素
        # with open("GA_result.csv", 'a', encoding='utf-8', newline='') as f:  # 打开文件并记录评估次数和最佳适应度
        #     write = csv.writer(f)  # 创建writer对象
        #     write.writerow(list([self.best_fitness, best_code]))  # 记录最佳适应度和基因序列
        return best_code, self.best_fitness  # 返回最佳适应度和基因序列
